#include <gflags/gflags.h>
#include <gtest/gtest.h>
#include <thread>

#include "etcd.hpp"
#include "channel.hpp"
#include "logger.hpp"
#include "utils.hpp"

#include "file.pb.h"
#include "base.pb.h"

DEFINE_bool(run_mode, true, "是否为调试模式");
DEFINE_string(log_file, "", "发布模式下，用于指定日志的输出文件");
DEFINE_int32(log_level, 0, "发布模式下，用于指定日志输出等级");

DEFINE_string(etcd_host, "http://127.0.0.1:2379", "服务注册中心地址");
DEFINE_string(base_service, "/service", "服务监控根目录");
DEFINE_string(file_service, "/service/file_service", "服务监控根目录");


// 测试用例的基类，用于初始化和清理测试环境
class FileClientTest : public ::testing::Test {
protected:
    void SetUp() override {
        // 初始化服务发现和 RPC 信道
        auto sm = std::make_shared<im::ServiceManager>();
        sm->declared(FLAGS_file_service);
        auto put_cb = std::bind(&im::ServiceManager::onServiceOnline, sm.get(), std::placeholders::_1, std::placeholders::_2);
        auto del_cb = std::bind(&im::ServiceManager::onServiceOffline, sm.get(), std::placeholders::_1, std::placeholders::_2);
        im::Discover::ptr dclient = std::make_shared<im::Discover>(FLAGS_etcd_host, FLAGS_base_service, put_cb, del_cb);
        channel = sm->choose(FLAGS_file_service);
        if (!channel) {
            im::LOG_CRITICAL("Failed to get RPC channel");
        }
    }

    void TearDown() override {
        // 清理资源
    }

    im::ServiceChannel::ChannelPtr channel;
};

// 测试上传单个文件
TEST_F(FileClientTest, TestPutSingleFileSuccess) {
    std::string body;
    ASSERT_TRUE(im::readFile("./Makefile", body));
    
    im::FileService_Stub stub(channel.get());

    im::PutSingleFileReq req;
    req.set_request_id("1111");
    req.mutable_file_data()->set_file_name("Makefile");
    req.mutable_file_data()->set_file_size(body.size());
    req.mutable_file_data()->set_file_content(body);

    brpc::Controller cntl;
    im::PutSingleFileRsp rsp;
    stub.PutSingleFile(&cntl, &req, &rsp, nullptr);

    ASSERT_FALSE(cntl.Failed());
    ASSERT_TRUE(rsp.success());
    ASSERT_EQ(rsp.file_info().file_size(), body.size());
    ASSERT_EQ(rsp.file_info().file_name(), "Makefile");
}

// 测试下载单个文件
TEST_F(FileClientTest, TestGetSingleFileSuccess) {
    // 上传文件，获取文件 ID
    std::string body;
    ASSERT_TRUE(im::readFile("./Makefile", body));
    im::FileService_Stub stub(channel.get());
    im::PutSingleFileReq req;
    req.set_request_id("1111");
    req.mutable_file_data()->set_file_name("Makefile");
    req.mutable_file_data()->set_file_size(body.size());
    req.mutable_file_data()->set_file_content(body);
    brpc::Controller cntl;
    im::PutSingleFileRsp rsp;
    stub.PutSingleFile(&cntl, &req, &rsp, nullptr);
    ASSERT_TRUE(rsp.success());
    std::string file_id = rsp.file_info().file_id();

    // 下载文件
    im::GetSingleFileReq get_req;
    get_req.set_request_id("2222");
    get_req.set_file_id(file_id);
    brpc::Controller get_cntl;
    im::GetSingleFileRsp get_rsp;
    stub.GetSingleFile(&get_cntl, &get_req, &get_rsp, nullptr);

    ASSERT_FALSE(get_cntl.Failed());
    ASSERT_TRUE(get_rsp.success());
    ASSERT_EQ(get_rsp.file_data().file_id(), file_id);
    ASSERT_EQ(get_rsp.file_data().file_content(), body);
}

// 测试上传多个文件
TEST_F(FileClientTest, TestPutMultiFileSuccess) {
    std::string body1, body2;
    ASSERT_TRUE(im::readFile("./base.pb.h", body1));
    ASSERT_TRUE(im::readFile("./file.pb.h", body2));

    im::FileService_Stub stub(channel.get());
    im::PutMultiFileReq req;
    req.set_request_id("3333");

    auto file_data1 = req.add_file_data();
    file_data1->set_file_name("base.pb.h");
    file_data1->set_file_size(body1.size());
    file_data1->set_file_content(body1);

    auto file_data2 = req.add_file_data();
    file_data2->set_file_name("file.pb.h");
    file_data2->set_file_size(body2.size());
    file_data2->set_file_content(body2);

    brpc::Controller cntl;
    im::PutMultiFileRsp rsp;
    stub.PutMultiFile(&cntl, &req, &rsp, nullptr);

    ASSERT_FALSE(cntl.Failed());
    ASSERT_TRUE(rsp.success());
    ASSERT_EQ(rsp.file_info_size(), 2);
    ASSERT_EQ(rsp.file_info(0).file_name(), "base.pb.h");
    ASSERT_EQ(rsp.file_info(1).file_name(), "file.pb.h");
}

// 测试下载多个文件
TEST_F(FileClientTest, TestGetMultiFileSuccess) {
    // 上传多个文件
    std::string body1, body2;
    ASSERT_TRUE(im::readFile("./base.pb.h", body1));
    ASSERT_TRUE(im::readFile("./file.pb.h", body2));
    im::FileService_Stub stub(channel.get());
    im::PutMultiFileReq req;
    req.set_request_id("3333");
    auto file_data1 = req.add_file_data();
    file_data1->set_file_name("base.pb.h");
    file_data1->set_file_size(body1.size());
    file_data1->set_file_content(body1);
    auto file_data2 = req.add_file_data();
    file_data2->set_file_name("file.pb.h");
    file_data2->set_file_size(body2.size());
    file_data2->set_file_content(body2);
    brpc::Controller cntl;
    im::PutMultiFileRsp rsp;
    stub.PutMultiFile(&cntl, &req, &rsp, nullptr);
    ASSERT_TRUE(rsp.success());
    ASSERT_EQ(rsp.file_info_size(), 2);
    std::vector<std::string> file_ids;
    for (int i = 0; i < rsp.file_info_size(); ++i) {
        file_ids.push_back(rsp.file_info(i).file_id());
    }

    // 下载多个文件
    im::GetMultiFileReq get_req;
    get_req.set_request_id("4444");
    for (const auto& id : file_ids) {
        get_req.add_file_id_list(id);
        im::LOG_DEBUG("----- file_id={}-------", id);
    }
    brpc::Controller get_cntl;
    im::GetMultiFileRsp get_rsp;
    stub.GetMultiFile(&get_cntl, &get_req, &get_rsp, nullptr);

    ASSERT_FALSE(get_cntl.Failed());
    EXPECT_TRUE(get_rsp.success());
    EXPECT_EQ(get_rsp.file_data().size(), file_ids.size());
    for (const auto& id : file_ids) {
        ASSERT_TRUE(get_rsp.file_data().find(id) != get_rsp.file_data().end());
    }
}

// 测试下载不存在的文件
TEST_F(FileClientTest, TestGetSingleFileNonExistent) {
    im::FileService_Stub stub(channel.get());
    im::GetSingleFileReq req;
    req.set_request_id("2222");
    req.set_file_id("non_existent_file_id");

    brpc::Controller cntl;
    im::GetSingleFileRsp rsp;
    stub.GetSingleFile(&cntl, &req, &rsp, nullptr);

    ASSERT_FALSE(cntl.Failed());
    ASSERT_FALSE(rsp.success()); // 预期下载失败
}

// 测试下载多个不存在的文件
TEST_F(FileClientTest, TestGetMultiFileNonExistent) {
    im::FileService_Stub stub(channel.get());
    im::GetMultiFileReq req;
    req.set_request_id("4444");
    req.add_file_id_list("non_existent_file1");
    req.add_file_id_list("non_existent_file2");

    brpc::Controller cntl;
    im::GetMultiFileRsp rsp;
    stub.GetMultiFile(&cntl, &req, &rsp, nullptr);

    ASSERT_FALSE(cntl.Failed());
    ASSERT_FALSE(rsp.success()); // 预期下载失败
}

int main(int argc, char *argv[]) {
    testing::InitGoogleTest(&argc, argv);
    google::ParseCommandLineFlags(&argc, &argv, true);
    im::initLogger(FLAGS_run_mode, FLAGS_log_file, FLAGS_log_level);

    return RUN_ALL_TESTS();
}